import sys
from pathlib import Path
import shutil
import bentoml


def _write_latest(tag):
    model = bentoml.models.get(tag)
    model_store_path = Path(model.path).parent
    latest_file = model_store_path / "latest"
    if latest_file.exists():
        if latest_file.is_dir():
            shutil.rmtree(latest_file, ignore_errors=True)
        else:
            latest_file.unlink()
    latest_file.write_text(model.tag.version)
    return model.tag.version


def main():
    args = sys.argv[1:]
    if args:
        version = args[0]
        target_tag = f"iris_clf:{version}"
    else:
        models = list(bentoml.models.list("iris_clf"))
        if not models:
            print("No iris_clf models found.")
            return 1
        models.sort(key=lambda m: m.creation_time, reverse=True)
        target_tag = str(models[0].tag)

    updated_version = _write_latest(target_tag)
    print(f"✓ Updated 'latest' tag to: {updated_version}")

    latest_model = bentoml.models.get("iris_clf:latest")
    latest_version = latest_model.tag.version
    models = list(bentoml.models.list("iris_clf"))
    models.sort(key=lambda m: m.creation_time, reverse=True)
    for m in models:
        marker = " <- latest" if m.tag.version == latest_version else ""
        print(f"{m.tag}{marker}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())

